import os
import h5py
import numpy as np
import scipy.io.wavfile as wav
from python_speech_features import mfcc, delta


PHONE61=['iy','ih','eh','ey','ae','aa','aw','ay','ah','ao',
       'oy','ow','uh','uw','ux','er','ax','ix','axr','ax-h'
       'jh','ch','b','d','g','p','t','k','dx','s',
       'sh','z','zh','f','th','v','dh','m','n','ng',
       'em','nx','en','eng','l','r','w','y','hh','hv',
       'el','bcl','dcl','gcl','pcl','tcl','kcl','q','pau','epi','h#']
PHONE48=['iy','ih','eh','ey','ae','aa','aw','ay','ah','ao',
         'oy','ow','uh','uw','er','ax','ix','jh','ch','b',
         'd','g','p','t','k','dx','s','sh','z','zh',
         'f','th','v','dh','m','n','ng','en','l','r',
         'w','y','hh','el','vcl','cl','epi','sil']
PHONE39=['iy','ih','eh','ey','ae','aa','aw','ay','ah','oy',
         'ow','uh','uw','er','jh','ch','b','d','g','p',
         't','k','dx','s','sh','z','f','th','v','dh',
        'm','n','ng','l','r','w','y','hh','cl']
dic={'ux':'uw','axr':'er','em':'m','nx':'n','eng':'ng','hv':'hh',
     'bcl':'cl','dcl':'cl','gcl':'cl','pcl':'cl','tcl':'cl',
     'kcl':'cl','qcl':'cl','vcl':'cl','epi':'cl','pau':'cl','h#':'cl','el':'l','en':'n',
     'zh':'sh','ao':'aa','ix':'ih','ax':'ah','ax-h':'ah','q':'cl'}


base_dataset_path = 'D:/Datasets/TIMIT/raw'
stack_path = 'D:/Datasets/TIMIT/stack'

MAX_LENGTH = 256


def get_audio_feature(file_name, anchor = None, length = None):
  '''
  获取wav文件提取mfcc特征之后的数据
  '''

  audio_filename = file_name

  #读取wav文件内容，fs为采样率， audio为数据
  fs, audio = wav.read(audio_filename)

  #提取mfcc特征
  inputs = mfcc(audio, samplerate=fs)
  d_mfcc_1 = delta(inputs, 1)
  d_mfcc_2 = delta(inputs, 2)
  feature_inputs = np.hstack((inputs, d_mfcc_1, d_mfcc_2))
  # 对特征数据进行归一化，减去均值除以方差
  # feature_inputs = np.asarray(inputs[np.newaxis, :])
  # feature_inputs = (feature_inputs - np.mean(feature_inputs))/np.std(feature_inputs)

  #特征数据的序列长度
  feature_seq_len = feature_inputs.shape[0]

  if anchor is None and length is None :
    return feature_inputs, feature_seq_len
  else:
    return feature_inputs[anchor : anchor + length], length


def get_audio_label(file_name, time_s = None, time_e = None, need_time = False):
  '''
  将phone label转换成整数序列
  '''
  target_filename = file_name

  with open(target_filename, 'r') as f:
    phone=[]
    targets=[]
    time_pairs=[]
    lines = f.readlines()
    for line in lines:
        splits = line.split(' ')
        time_pairs.append((int(splits[0]), int(splits[1])))
        phone.append(splits[-1].split('\n')[0])
    for p in phone:
        if p in PHONE39:
            targets.append(PHONE39.index(p))
        elif dic[p] in PHONE39:
            targets.append(PHONE39.index(dic[p]))
        else:
            print('error')

    if time_s is not None and time_e is not None:
        s = 0
        e = 0
        for i in range(len(time_pairs)):
            if time_pairs[i][0] <= time_s <= time_pairs[i][1]:
            	front = time_s - time_pairs[i][0]
            	back = time_pairs[i][1] - time_s
            	if front <= back:
            		s = i
            	else:
            		s = i + 1
            if time_pairs[i][0] <= time_e <= time_pairs[i][1]:
            	front = time_e - time_pairs[i][0]
            	back = time_pairs[i][1] - time_e
            	if front <= back:
            		e = i
            	else:
            		e = i + 1

        if e <= s:
            e = len(time_pairs) - 1
        if s >= len(time_pairs):
            s = len(time_pairs) - 1

        targets = targets[s:e+1]
        time_pairs = time_pairs[s:e+1]

  return targets, time_pairs


def get_featureandlabel(mode):

    count1,count2=0,0
    dataset_path=os.path.join(base_dataset_path,mode)
    
    if os.path.exists(dataset_path) == False:
        return False

    lst_contents = os.listdir(dataset_path)
    fi=[]
    fs=[]
    l=[]
    times = []

    for k in range(0, len(lst_contents)):
        second_path=os.path.join(dataset_path,lst_contents[k])
        if os.path.isdir(second_path) == True:
            second_contents=os.listdir(second_path)
            for p in range(0,len(second_contents)):
                last_path=os.path.join(second_path,second_contents[p])
                if os.path.isdir(last_path) == True:
                    last_contents = os.listdir(last_path)
                    for w in range(0,len(last_contents)):
                        if last_contents[w].split('.')[-1]=='wav':
                            file_name = os.path.join(last_path, last_contents[w])

                            feature_inputs, feature_seq_len=get_audio_feature(file_name)
                            if feature_seq_len > MAX_LENGTH and 'TEST' in file_name:
                                n_clips = int(np.ceil(feature_seq_len / MAX_LENGTH))
                                n_overlap = int(np.floor((n_clips * MAX_LENGTH - feature_seq_len) / (n_clips - 1)))
                                for i in range(n_clips):
                                    anchor = i * (MAX_LENGTH - n_overlap)
                                    time_s = anchor * 160
                                    time_e = (anchor + MAX_LENGTH) * 160

                                    feature_inputs, feature_seq_len=get_audio_feature(file_name, anchor, MAX_LENGTH)
                                    fi.append(feature_inputs)
                                    fs.append(feature_seq_len)

                                    count1 += 1
                                    print(mode+' feature file %d' % count1 + ' has done.')

                                    label_name = file_name.split('_')[0] + '.PHN'
                                    label, _ =get_audio_label(label_name, time_s, time_e)
                                    l.append(label)

                                    count2 += 1
                                    print(mode+' label file %d'%count2+' has done.')
                            else:
                                fi.append(feature_inputs)
                                fs.append(feature_seq_len)

                                count1 += 1
                                print(mode+' feature file %d' % count1 + ' has done.')
                            
                                label_name = file_name.split('_')[0] + '.PHN'
                                if 'TRAIN' in file_name:
                                    label, time =get_audio_label(label_name, need_time = True)
                                    times.append(time)
                                else:
                                    label, _ =get_audio_label(label_name)
                                l.append(label)

                                count2 += 1
                                print(mode+' label file %d'%count2+' has done.')

    return fi,fs,l,times


def main():
    if os.path.exists(stack_path) == False:
        os.mkdir(stack_path)

    fi,fs,label,times=get_featureandlabel('TRAIN')
    os.chdir(stack_path)
    np.savez_compressed('train_phn',fi=fi,fs=fs,label=label,times=times)


    fi,fs,label,_=get_featureandlabel('TEST')
    os.chdir(stack_path)
    np.savez_compressed('test_phn', fi=fi, fs=fs,label=label)

    print("prepare data file is saved.")


if __name__ == '__main__':
    main()
